# roombafollowtheleader.py
# Carl Nygaard
# Max Pfleuger
# Kyle Roberts

# April 2007
# HMC Robotics

import roomba
import msvcrt
import time
import math
import sys
import threading
import random
import time

# Usage:  
# python filename.py leaderPortNum follower1PortNum follower2PortNum ...


if len(sys.argv) > 1:
    # Load in a list of ports
    PORT = [int(val) for val in sys.argv[1:]]
else:
    PORT = 10


class MarchingRoombaThread(threading.Thread):
    """A convenient absract Roomba thread class."""

    def __init__(self, port):
        threading.Thread.__init__(self)
        self.port = port
        self.running = True
        self.isVirtualWallTriggered = False
        self.follower = None
        # Generic States are Negative
        self.WAITING = -1
        # Previous state of roomba
        self.prevState = None

    def run(self):
        self.reset()

        # Main loop
        while self.running:
            # Update whether or not we sense the virtual wall
            self.updateVirtualWall()

            #Check if the follower is waiting
            self.checkFollower()
            
            # Perform specific robot tasks
            self.step()
                      

            # Take a breather
            time.sleep(0.15)

        # Exit gracefully and stop robot
        self.thebot.roomba.drive(0,0)
        self.thebot.roomba.ser.close()

    def step(self):
        pass

    def stop(self):
        self.running = False

    def reset(self):
        self.thebot = roomba.RoombaRobot(self.port)
        time.sleep(.2)
        self.sendSequence([128, 131, 128, 131])

    def setFollower(self, follower):
        self.follower = follower

    def sendSequence(self, sequence):
        for val in sequence:
            self.thebot.roomba.ser.write(chr(val))

    def requestSensors(self, sensors):
        """Request sensors from roomba"""
        self.sendSequence([149, len(sensors)] + sensors)
        r = self.thebot.roomba.ser.read(size = len(sensors))
        return [ ord(c) for c in r]
     
    def updateVirtualWall(self):
        """Returns True if the virtual wall is sensed"""
        # Request the virtual wall
        sensorData = self.requestSensors([13])
        self.isVirtualWallTriggered = sensorData[0]
        
        # Light LED if detecting IR
        if self.isVirtualWallTriggered:
            self.sendSequence([139, 2, 128, 255])
        else:
            self.sendSequence([139, 0, 0, 0])
    
    def getBumpStatus(self):
        """Returns isLeftBumping, isRightBumping"""
        sensorData = self.requestSensors([7])
        leftBump = roomba.bitOfByte( 1, sensorData[0])
        rightBump = roomba.bitOfByte(0, sensorData[0])
        return leftBump, rightBump

    def turn(self, turnAngle):
        """Turn the robot turnAngle number of degrees"""

        if (turnAngle < 0):
            high, low = roomba.toTwosComplement2Bytes(-300)
        else:
            high, low = roomba.toTwosComplement2Bytes(300)

        high1, low1 = roomba.toTwosComplement2Bytes(int(turnAngle))
        turnScript = [
                 137, #Drive
                 high, low,
                 # Turn Clockwise
                 0, 0,
                 # Wait for angle
                 157,
                 high1, low1,
                 # Stop
                 137, 0, 0, 0, 0,
                 # Send a sensor reading
                 149, 1, 13]

        # Upload the script
        self.sendSequence(
                [152,              # Prepare to receive script
                 len(turnScript)]) # Length of script
        self.sendSequence(turnScript)
        # Play the script
        self.sendSequence([153])

        # Wait for the turn to finish (which happens when the sensor reading is
        # sent back to the computer
        start = time.time()
        time.sleep(0.1)
        r = []
        # Wait until we get a value back
        while len(r) == 0:
            r = self.thebot.roomba.ser.read(size = 1)
        print "Finished turning after seconds:", time.time() - start

    def mouseDemo(self):
        #Start mouse demo
        self.sendSequence([136, 3])

    def endDemo(self):
        #Stop current demo, Put roomba back in full mode
        self.sendSequence([136, 255, 132])
    def checkFollower(self):
        #Make sure follower has a signal
        if self.follower != None:
            noLeader = (self.follower.state == self.follower.TURNINGRIGHT or 
                self.follower.state == self.follower.TURNINGLEFT  or
                self.follower.state == self.follower.WAITING)
            if (noLeader):
                if (self.state != self.WAITING):
                    self.prevState = self.state
                    self.state = self.WAITING
            elif (self.state == self.WAITING):
                self.state = self.prevState

class LeaderRoombaThread(MarchingRoombaThread):
    """The leader of the pack.  Only one of these should be created."""

    def __init__(self, port):
        MarchingRoombaThread.__init__(self, port)

        # Speed of user control
        self.speed = 200

        # States
        self.USERCONTROLLED = 0
        self.BOUNCING = 1
        self.FINDFOLLOWER = 2

        # Rediscover off to begin with
        self.rediscover = False

        # The initial/current state
        self.state = self.USERCONTROLLED
        # The previous state
        self.prevState = self.USERCONTROLLED

    def step(self):
        """Performs a single robot iteration."""
        # Handle user input
        self.processKeyboard()
       
        """# If we aren't being followed backup and disable user control
        if self.rediscover == True:
            if not self.isVirtualWallTriggered:
                if self.state != self.FINDFOLLOWER:
                    self.prevState = self.state
                self.state = self.FINDFOLLOWER
            elif self.state == self.FINDFOLLOWER:
                self.state = self.prevState
                self.thebot.roomba.drive(0,0) """
        
        global aRoombaIsLost
        aRoombaIsLost = (self.state == self.WAITING)

        # DO STATE SPECIFIC BEHAVIOR
        # Perform bouncing, if enabled
        if self.state == self.BOUNCING:
            self.bounceStep()
        elif self.state == self.FINDFOLLOWER:
            self.thebot.roomba.drive(-75, 2001)
        elif self.state == self.WAITING:
            self.thebot.roomba.drive(0,0)

    def bounceStep(self):
        """Perform one iteration in bounce mode. May include a full 
           turn step."""
        leftBump, rightBump = self.getBumpStatus()
        if leftBump or rightBump:
            # Turn left or right some amount greater than 60 degrees
            if rightBump:
                turnAmount = 60 + 90 * random.random()
            else:
                turnAmount = -60 - 90 * random.random()
            self.turn(turnAmount)
        # Drive forward
        self.thebot.roomba.drive(225,2001)

    def processKeyboard(self):
        if msvcrt.kbhit():
            c = msvcrt.getch()
            # Start and stop commands
            if c == 'b':
                print "Starting to bounce around!"
                self.state = self.BOUNCING
            if c == 'r':
                print "Toggle Rediscover"
                self.rediscover = not self.rediscover
            if c == 'w' or c == ' ':
                print 'Stopping'
                self.state = self.USERCONTROLLED
                self.thebot.roomba.drive(0,0)
            if c == '+':
                self.speed += 10
                print 'User Control Speed is now: ', self.speed
            if c == '-':
                self.speed -= 10
                print 'User Control Speed is now: ', self.speed 
            if c == 'q':
                stopAllThreads()

            # Manual control
            if self.state == self.USERCONTROLLED:
                if c == 'i':
                    print 'i'
                    self.thebot.roomba.drive(self.speed,2001)
                if c == 'k':
                    print 'k'
                    self.thebot.roomba.drive(-1*self.speed / 2, 2001)
                if c == 'j':
                    print 'j'
                    self.thebot.roomba.drive(self.speed, 300)
                if c == 'l':
                    print 'l'
                    self.thebot.roomba.drive(self.speed, -300)
                if c == 'u':
                    print 'u'
                    self.thebot.roomba.drive(self.speed,0)

class FollowerRoombaThread(MarchingRoombaThread):
    """The following roomba."""

    def __init__(self, port, leaderRobot):
        MarchingRoombaThread.__init__(self, port)
    
        # The robot this robot is following
        self.leaderRobot = leaderRobot

        # States
        # JUSTTURNED<dir> indicates which way the roomba last turned to try to
        # face the leader.  TURNGING<dir> indicates that we are in the process
        # of turning a given direction.
        self.JUSTTURNEDRIGHT = 0
        self.JUSTTURNEDLEFT = 1
        self.TURNINGRIGHT = 2
        self.TURNINGLEFT = 3
        # ZIGLEFT and ZAGRIGHT will keep us oscillating on the leader when following.
        # In most cases this should allow you to know which way to turn to find the leader.
        self.ZIGLEFT = 4
        self.ZAGRIGHT = 5

        # The initial state
        self.state = self.JUSTTURNEDRIGHT
        self.prevState = self.JUSTTURNEDRIGHT
        self.leaderFound = False # this is to make Zig-Zag work
        self.zigZagTime = time.time() # keeps time spent searching for the leader
        self.speedFactor = 1 #speed factor to slow down if roomba is lost

    def step(self):
        """Performs a single robot iteration."""
        global aRoombaIsLost
        if aRoombaIsLost:
            self.speedFactor = .1
        else:
            self.speedFactor = 1
        
        # Perform bouncing, if enabled
        if self.leaderRobot.isVirtualWallTriggered:
            # Update state information
            if self.state == self.TURNINGLEFT:
                #self.state = self.JUSTTURNEDLEFT
                self.state = self.ZIGLEFT
            elif self.state == self.TURNINGRIGHT:
                #self.state = self.JUSTTURNEDRIGHT
                self.state = self.ZAGRIGHT
            if (self.state == self.ZIGLEFT or self.state == self.ZAGRIGHT):
                self.leaderFound = True
                self.zigZagTime = time.time()
            
            # Get bump info
            [leftBump, rightBump] = self.getBumpStatus()
            bump = leftBump or rightBump
            
            # My leader sees me, drive forward
            if ((self.state == self.JUSTTURNEDLEFT 
                or self.state == self.JUSTTURNEDRIGHT) 
                and not bump):
                self.thebot.roomba.drive(200,2001)
            elif self.state == self.ZIGLEFT and not bump:
                self.thebot.roomba.drive(int(300* self.speedFactor), 250)
            elif self.state == self.ZAGRIGHT and not bump:
                self.thebot.roomba.drive(int(300* self.speedFactor), -250)
            
            if bump:
                self.thebot.roomba.drive(-100,2001)
            
        else:
            # Update state information
            if self.state == self.JUSTTURNEDRIGHT:
                self.state = self.TURNINGLEFT
            elif self.state == self.JUSTTURNEDLEFT:
                self.state = self.TURNINGRIGHT
            elif self.leaderFound == True:
                self.leaderFound = False
                if self.state == self.ZIGLEFT:
                    self.state = self.ZAGRIGHT
                elif self.state == self.ZAGRIGHT:
                    self.state = self.ZIGLEFT
            
            # If we can't find the leader we may need to change algorithms
            if (self.state == self.ZIGLEFT or self.state == self.ZAGRIGHT) and self.leaderFound == False:
                    if time.time() - self.zigZagTime > 1.5:
                        if self.state == self.ZIGLEFT:
                            self.state = self.TURNINGLEFT
                        elif self.state == self.ZAGRIGHT:
                            self.state = self.TURNINGRIGHT
            
            # Get bump info
            [leftBump, rightBump] = self.getBumpStatus()
            bump = leftBump or rightBump
            
            # My leader does not see me, turn
            if self.state == self.TURNINGLEFT and not bump:
                self.thebot.roomba.drive(200,0)
            elif self.state == self.TURNINGRIGHT and not bump:
                self.thebot.roomba.drive(-200,0)
            elif self.state == self.ZIGLEFT and not bump:
                    self.thebot.roomba.drive(int(300* self.speedFactor), 250)
            elif self.state == self.ZAGRIGHT and not bump:
                self.thebot.roomba.drive(int(300* self.speedFactor), -250)
 

            if bump:
                self.thebot.roomba.drive(-100,2001)
                print 'bumping'
       


        if self.state == self.WAITING:
            self.thebot.roomba.drive(0,0)     

def stopAllThreads():
    for thread in threads:
        thread.stop()

# Holds all the different threads, in the order of leader, follower, follower,
# etc...
threads = []
aRoombaIsLost = False


if __name__ == '__main__':      # are we at the command line
    threads.append(LeaderRoombaThread(PORT[0]))
    for port in PORT[1:]:
        print 'starting up port', port
        threads.append(FollowerRoombaThread(port, threads[-1]))
        threads[-2].setFollower(threads[-1])

    # Start rombas in reverse order
    for thread in threads[::-1]:
        thread.start()
